'''
Created on Feb 9, 2012

@author: xzhu
'''
import sys
import ioniState, stateHead

class AllState():
    '''
    All the protonation states
    '''
    # sHead is a reference to an existing KeyResidues object.
    sHead = stateHead.KeyResidues()
    
    def __init__(self):
        '''
        Constructor
        '''
        self.states = []
        
        
    def load_all_states(self, lines):
        '''
        Load all the possible ionization states from lines of a file
        
        By default, this file is named "ener.txt".
        '''
        count = 0
        for line in lines:
            buff_state = ioniState.IoniState()
            buff_state.load_state(line)
      
            buff_state.seq = count
            count += 1
            self.states.append(buff_state)
        
    def getAllPossibleStates(self):
        import itertools
        for crgs in itertools.product(*AllState.sHead.possibleChargesOfKeyResidues):
            new_state = ioniState.IoniState()
            new_state.crgs = list(crgs)
            self.states.append(new_state)
                   
    def state_of_crg(self, crgs):
        '''
        return the state which has the same charges as crgs
        '''
        #print 'len  ', len(self.states)
        for state in self.states:
            if state.crgs == crgs:
                return state
           
        print >> sys.stderr, 'can not find the crgs', crgs
        
            
    def write_seq_ID(self, outfile = 'allStates_seq_ID.txt'):
        '''
        output the seq and ID relationship of all the states
        '''
        
        ofile = open(outfile, 'w')
        for state in self.states:
            print state.ID, state.seq
            ofile.write(str(state.seq) + '\t' + state.ID + '\n')
        ofile.close()
        